Add ListMLE Loss #130

LakshmiKalaKadali · 2025-06-24T16:33:33Z

Added ListMLELoss code to listwise ranking. This code does not consider Lambda weights.
Here is the gist for verified results with TFRS ListMLELoss

abheesht17 · 2025-08-07T02:58:56Z

Not sure why tests did not run, commenting so that tests run

abheesht17

Sorry for the delay, did one pass! Some questions below, let me know what you think.

I'll take a closer look again!

api_gen.py

abheesht17 · 2025-08-07T02:59:49Z

keras_rs/api/layers/__init__.py

+from keras_rs.src.layers.embedding.distributed_embedding import DistributedEmbedding as DistributedEmbedding
+from keras_rs.src.layers.embedding.distributed_embedding_config import FeatureConfig as FeatureConfig
+from keras_rs.src.layers.embedding.distributed_embedding_config import TableConfig as TableConfig
+from keras_rs.src.layers.embedding.embed_reduce import EmbedReduce as EmbedReduce
+from keras_rs.src.layers.feature_interaction.dot_interaction import DotInteraction as DotInteraction
+from keras_rs.src.layers.feature_interaction.feature_cross import FeatureCross as FeatureCross
+from keras_rs.src.layers.retrieval.brute_force_retrieval import BruteForceRetrieval as BruteForceRetrieval
+from keras_rs.src.layers.retrieval.hard_negative_mining import HardNegativeMining as HardNegativeMining
+from keras_rs.src.layers.retrieval.remove_accidental_hits import RemoveAccidentalHits as RemoveAccidentalHits


Hmmm, these shouldn't be getting re-formatted, in my opinion. Are you on the correct Ruff version? Same for other files

keras_rs/src/losses/list_mle_loss_test.py

keras_rs/src/losses/list_mle_loss.py

LakshmiKalaKadali · 2025-09-15T16:14:15Z

Hi @abheesht17, Done suggested changes. Could you please review the PR. Thank you

abheesht17

This LGTM! I'm going to drop an approval on this, but could you please format the code first?

abheesht17 · 2025-10-08T01:23:05Z

Not sure why we don't have a way to run tests here, need to figure that out before merging

abheesht17

Noticed a few nits. Instead of using +, -, *, /, let's use Keras ops everywhere

keras_rs/src/losses/list_mle_loss.py

abheesht17 · 2025-10-13T14:23:39Z

@LakshmiKalaKadali - can you look into the test failures above?

LakshmiKalaKadali · 2025-10-16T07:15:33Z

Hi @abheesht17,

The torch test fails are due to precision mismatch in between actual and expected values. That precision mismatch might be due to the mask AFTER exp(). Applied the mask to set invalid positions to (-1e9) before calling ops.exp(), and removed the masking that happens after exp(). I hope that this resolves the torch test failure. Could you please run the tests. Thank You

hertschuh · 2025-10-22T17:42:21Z

@LakshmiKalaKadali

The version of torch was updated for the tests. Please rebase to see if it helps.

hertschuh · 2025-10-23T21:50:45Z

@LakshmiKalaKadali

I do not understand why the results are different. But the differences you see are significant enough that it has to be a bug. Can you try to debug by printing all intermediate values in compute_unreduced_loss and call and comparing the values for Torch and JAX or TF? Hopefully we can narrow it down to a single line of code.

LakshmiKalaKadali · 2025-10-24T04:30:38Z

Sure @hertschuh, I will deep dive into the intermediate values. Thank You

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

LakshmiKalaKadali · 2025-10-29T16:36:02Z

Hi @hertschuh, all the tests are passing now, Could you please take a look. Thank You

hertschuh

Thanks for figuring out the issue!

hertschuh · 2025-10-30T02:14:27Z

.github/workflows/actions.yml

    - name: Test with pytest
      run: |
-        pytest keras_rs/
+        pytest keras_rs/ 


Undo this file.

hertschuh · 2025-10-30T02:19:12Z

keras_rs/src/losses/list_mle_loss.py

+        # added stable offset before calling sort_by_scores
+        list_size = ops.shape(labels_for_sorting)[1]
+        indices = ops.arange(list_size)
+
+        indices = ops.expand_dims(indices, axis=0)
+        indices = ops.broadcast_to(indices, ops.shape(labels_for_sorting))
+
+        stable_offset = ops.cast(indices, labels_for_sorting.dtype) * 1e-6
+
+        labels_for_sorting = ops.subtract(labels_for_sorting, stable_offset)


Oh... do the issue is that Torch's topk is not stable with sorted=True, unlike JAX and TF.

There is not easy fix for topk that would be efficient.

However, do you mind:

moving this in the sort_by_scores in ranking_metrics_utils.py

doing it only if keras.backend.backend() == "torch" and shuffle_ties is False

Thanks!

Can you add a comment with a reference to these:

Make topk sort stable pytorch/pytorch#27542

Enable torch.topk to support stable flag pytorch/pytorch#88227

We will be able to remove this code if they ever fix the "stable" issue.

hertschuh · 2025-10-30T15:57:11Z

keras_rs/src/metrics/ranking_metrics_utils.py

+    # equal scores. We can remove this workaround once PyTorch adds a
+    # `stable=True` flag for topk.
+
+    if K.backend() == "torch" and not shuffle_ties:


nitpick: K was how keras was imported many many years ago. But we don't use that style anymore. Simply remove the import and write keras.backend.backend()

hertschuh

Thanks!

LakshmiKalaKadali marked this pull request as ready for review June 25, 2025 16:37

laxmareddyp requested review from abheesht17 June 26, 2025 15:22

abheesht17 requested changes Aug 11, 2025

View reviewed changes

LakshmiKalaKadali force-pushed the listmleloss branch from ef5ee24 to 63388e9 Compare August 18, 2025 10:16

abheesht17 reviewed Oct 8, 2025

View reviewed changes

abheesht17 approved these changes Oct 8, 2025

View reviewed changes

abheesht17 reviewed Oct 8, 2025

View reviewed changes

keras_rs/src/losses/list_mle_loss.py Outdated Show resolved Hide resolved

keras_rs/src/losses/list_mle_loss.py Outdated Show resolved Hide resolved

keras_rs/src/losses/list_mle_loss.py Outdated Show resolved Hide resolved

LakshmiKalaKadali force-pushed the listmleloss branch 2 times, most recently from e986ef5 to 3da3fa2 Compare October 23, 2025 07:43

LakshmiKalaKadali and others added 10 commits October 24, 2025 13:21

Add STU layer

730bc21

Update keras_rs/src/layers/stu.py

1906f2d

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Update keras_rs/src/layers/common.py

d0add4f

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Update keras_rs/src/layers/jagged_tensors.py

eb98ae1

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Update keras_rs/src/layers/hstu_mha_attention.py

0347b51

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Update keras_rs/src/layers/stu.py

570df2b

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Update keras_rs/src/layers/stu.py

6f3dcef

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Update keras_rs/src/layers/jagged_tensors.py

c80e598

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Update keras_rs/src/layers/hstu_compute_output.py

8073dcb

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Add list mle loss

8cb2d98

LakshmiKalaKadali force-pushed the listmleloss branch from 157f774 to 8cb2d98 Compare October 24, 2025 08:11

Debug statements added

b643778

LakshmiKalaKadali added 4 commits October 29, 2025 18:13

pytest added

3f6418c

Added_stable offset code to sorting_labels

47f0251

Stable_offset code to clear the ambiguity with similar labels

a3679c6

Stable_offset code added to handle sorting_label

22757fd

hertschuh reviewed Oct 30, 2025

View reviewed changes

LakshmiKalaKadali force-pushed the listmleloss branch from 0761a46 to 22757fd Compare October 30, 2025 07:18

LakshmiKalaKadali and others added 4 commits October 30, 2025 12:53

Merge branch 'keras-team:main' into listmleloss

7564eb5

Save local changes

35bfa0d

Updated few code changes

8424416

lint errors corrected

7017f35

hertschuh reviewed Oct 30, 2025

View reviewed changes

deleted keras backend import statement

7044610

hertschuh approved these changes Oct 30, 2025

View reviewed changes

hertschuh merged commit 04fb241 into keras-team:main Oct 30, 2025
7 checks passed

Add ListMLE Loss #130

Add ListMLE Loss #130

Uh oh!

Conversation

LakshmiKalaKadali commented Jun 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

abheesht17 commented Aug 7, 2025

Uh oh!

abheesht17 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

abheesht17 Aug 7, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

LakshmiKalaKadali commented Sep 15, 2025

Uh oh!

abheesht17 left a comment

Choose a reason for hiding this comment

Uh oh!

abheesht17 commented Oct 8, 2025

Uh oh!

abheesht17 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

abheesht17 commented Oct 13, 2025

Uh oh!

LakshmiKalaKadali commented Oct 16, 2025

Uh oh!

hertschuh commented Oct 22, 2025

Uh oh!

hertschuh commented Oct 23, 2025

Uh oh!

LakshmiKalaKadali commented Oct 24, 2025

Uh oh!

LakshmiKalaKadali commented Oct 29, 2025

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

hertschuh Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

hertschuh Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

hertschuh Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

hertschuh Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

LakshmiKalaKadali commented Jun 24, 2025 •

edited

Loading